package com.vilen.cms.util;


import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.xml.sax.InputSource;

import java.io.*;
import java.sql.SQLException;

/**
 * Created by vilen on 2015/8/24.
 */
public class AbstractDbnitTestCase {
    public static IDatabaseConnection dbunitCon;
    private File tempFile;
    private File dtdFile;

    @BeforeClass
    public static void init() throws DatabaseUnitException, SQLException {
        dbunitCon = new DatabaseConnection(DbUtil.getConnection());
    }

    protected IDataSet createDateSet(String tname) throws DataSetException, IOException {
        InputStream is = AbstractDbnitTestCase.class.getClassLoader().getResourceAsStream(tname + ".xml");
        Assert.assertNotNull("dbunit的基本数据文件不存在", is);
        return new FlatXmlDataSet(new FlatXmlProducer(new InputSource(is), new FlatDtdDataSet(new FileReader(dtdFile))));
    }

    private void writeBackupFile(IDataSet ds) throws IOException, DataSetException {
        tempFile = File.createTempFile("back","xml");
        dtdFile = File.createTempFile("back", "dtd");
//        tempFile = new File("D:/test.xml");
//        dtdFile = new File("D:/test.dtd");
        FlatDtdDataSet.write(ds, new FileWriter(dtdFile));
        FlatXmlDataSet.write(ds, new FileWriter(tempFile));
    }

    protected void backupAllTable() throws SQLException, IOException, DataSetException {
        IDataSet ds = dbunitCon.createDataSet();
        writeBackupFile(ds);
    }

    protected void backupCustomTable(String[] tname) throws DataSetException, IOException {
        QueryDataSet ds = new QueryDataSet(dbunitCon);
        for (String str : tname) {
            ds.addTable(str);
        }
        writeBackupFile(ds);
    }

    private void backupOneTable(String tname) throws DataSetException, IOException {
        backupCustomTable(new String[]{tname});
    }

    protected void resumeTable() throws IOException, DatabaseUnitException, SQLException {
        IDataSet ds = new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(tempFile)), new FlatDtdDataSet(new FileReader(dtdFile))));
        DatabaseOperation.CLEAN_INSERT.execute(dbunitCon, ds);
    }

    @AfterClass
    public static void destory() {
        try {
            if (dbunitCon != null) {
                dbunitCon.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }
}
